#单眼
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
from model.model import MultiLabelBiFPNResNet
import os

# 定义类别标签（顺序与模型输出一致）
def load_model(model_path=None):
    """加载训练好的模型"""
    if model_path is None:
        # 使用绝对路径加载模型
        model_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "best_resnet_model.pth")
    
    model = MultiLabelBiFPNResNet(num_classes=8)
    
    # 加载模型权重
    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
    
    # 处理前缀差异
    new_state_dict = {}
    for k, v in state_dict.items():
        # 将 'resnet' 前缀替换为 'backbone'
        if k.startswith('resnet.'):
            new_key = k.replace('resnet.', 'backbone.')
            new_state_dict[new_key] = v
        else:
            new_state_dict[k] = v
    
    # 加载模型
    model.load_state_dict(new_state_dict, strict=False)
    model.eval()
    return model


def preprocess_image(image_path):
    """加载并预处理图像"""
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    image = Image.open(image_path).convert("RGB")
    return transform(image).unsqueeze(0)  # 增加 batch 维度


def predict_both_eyes(left_image_path, right_image_path, threshold=0.5, model_path=None):
    """
    预测左右眼的疾病情况
    
    参数:
        left_image_path (str): 左眼图像路径
        right_image_path (str): 右眼图像路径
        threshold (float): 预测阈值，默认为0.5
        model_path (str): 模型路径，默认为None，会自动查找
        
    返回:
        dict: 包含预测结果的字典，包括各类别概率、预测的疾病和最终诊断
    """
    try:
        # 定义标签和映射
        labels = ['N', 'D', 'G', 'C', 'A', 'H', 'M', 'O']
        label_mapping = {
            'N': '正常', 'D': '糖尿病', 'G': '青光眼', 'C': '白内障',
            'A': 'AMD', 'H': '高血压', 'M': '近视', 'O': '其他疾病/异常'
        }
        
        # 加载模型
        model = load_model(model_path)
        
        # 预处理图像
        left_tensor = preprocess_image(left_image_path)
        right_tensor = preprocess_image(right_image_path)
        
        # 预测
        with torch.no_grad():
            outputs = model(left_tensor, right_tensor)
            probabilities = torch.sigmoid(outputs).squeeze().cpu().numpy()
        
        # 根据阈值确定预测结果
        predicted = (probabilities >= threshold).astype(int)
        
        # 获取预测的疾病（跳过正常标签N，只关注疾病）
        predicted_diseases = [labels[i] for i, pred in enumerate(predicted) if i > 0 and pred == 1]
        
        # 确定最终诊断
        if not predicted_diseases:
            diagnosis = label_mapping['N']
        else:
            diagnosis = "、".join([label_mapping[d] for d in predicted_diseases])
        
        # 构建结果字典
        result = {
            "probabilities": {label_mapping[labels[i]]: float(probabilities[i]) for i in range(len(labels))},
            "predicted_diseases": [label_mapping[d] for d in predicted_diseases],
            "diagnosis": diagnosis
        }
        
        return result
        
    except Exception as e:
        return {"error": str(e)}


def predict(model, left_image_path, right_image_path):
    """预测左右眼的疾病情况"""
    labels = ['N', 'D', 'G', 'C', 'A', 'H', 'M', 'O']
    label_mapping = {
        'N': '正常', 'D': '糖尿病', 'G': '青光眼', 'C': '白内障',
        'A': 'AMD', 'H': '高血压', 'M': '近视', 'O': '其他疾病/异常'
    }

    left_tensor = preprocess_image(left_image_path)
    right_tensor = preprocess_image(right_image_path)

    with torch.no_grad():
        outputs = model(left_tensor, right_tensor)
        probabilities = torch.sigmoid(outputs).squeeze().cpu().numpy()

    threshold = 0.5
    predicted = (probabilities >= threshold).astype(int)

    predicted_diseases = [labels[i] for i, pred in enumerate(predicted) if i > 0 and pred == 1]
    diagnosis = "、".join([label_mapping[d] for d in predicted_diseases]) if predicted_diseases else label_mapping['N']

    print("各类别概率:")
    for i, label in enumerate(labels):
        print(f"类别 {label_mapping[label]}: 概率 {probabilities[i]:.4f}")

    print("\n最终诊断结果:")
    print(diagnosis)


def main():
    """主函数"""
    model = load_model()
    left_image_path = "../Training_Dataset/14_left.jpg"
    right_image_path = "../Training_Dataset/14_right.jpg"
    predict(model, left_image_path, right_image_path)

if __name__ == "__main__":
    main()

